#!/bin/bash

# --- 默认配置 ---
MYSQL_USER="root"
MYSQL_PWD="test@2025"
SYSBENCH_EXECUTABLE="/usr/local/bin/sysbench"
LUA_SCRIPT_BASE_PATH="/root/sysbench_auto_opt/sysbench-1.0.20/src/lua" 
DB_NAME_PREFIX="sysbench"
SYSBENCH_OUTPUT_DIR="./sysbench_results"

# --- 默认的 Sysbench 参数 (如果未从命令行提供) ---
TABLE_SIZE=5000000
TABLE_COUNT=32
THREADS_PREPARE=64
THREADS_RUN=256
TEST_DURATION=120
REPORT_INTERVAL=10

# --- 函数：显示用法 ---
usage() {
    echo "用法: $0 -p <端口列表> -S <场景列表> [--prepare] [测试参数...]"
    echo "参数:"
    echo "  -p <列表>      用逗号分隔的 MySQL 实例端口列表 (必需)。"
    echo "  -S <列表>      用逗号分隔的 Sysbench 场景列表 (必需)。"
    echo "  --prepare      (标志) 运行 sysbench prepare 阶段。"
    echo "测试参数 (可选):"
    echo "  --table-size=<值>   覆盖默认的 TABLE_SIZE ($TABLE_SIZE)"
    echo "  --tables=<值>       覆盖默认的 TABLE_COUNT ($TABLE_COUNT)"
    echo "  --threads=<值>      覆盖默认的 THREADS_RUN ($THREADS_RUN)"
    echo ""
    exit 1
}

# --- 解析参数 ---
PREPARE_DATA="false"
PORTS_STR=""
SCENARIOS_STR=""

# 解析所有参数，包括 --key=value 格式的
while [[ "$#" -gt 0 ]]; do
    case $1 in
        -p) PORTS_STR="$2"; shift ;;
        -S) SCENARIOS_STR="$2"; shift ;;
        --prepare) PREPARE_DATA="true" ;;
        --table-size=*) TABLE_SIZE="${1#*=}" ;;
        --tables=*) TABLE_COUNT="${1#*=}" ;;
        --threads=*) THREADS_RUN="${1#*=}" ;;
        *) echo "未知参数: $1"; usage ;;
    esac
    shift
done


if [ -z "$PORTS_STR" ] || [ -z "$SCENARIOS_STR" ]; then
    usage
fi

IFS=',' read -r -a PORTS <<< "$PORTS_STR"
IFS=',' read -r -a SCENARIOS <<< "$SCENARIOS_STR"

# --- 函数：为单个实例运行 Sysbench ---
run_sysbench_for_instance() {
    local port=$1
    local db_name="${DB_NAME_PREFIX}_${TABLE_SIZE}x${TABLE_COUNT}_p${port}"
    local output_dir="${SYSBENCH_OUTPUT_DIR}/port_${port}"
    mkdir -p "$output_dir"

    echo "--- [端口 ${port}] 开始 Sysbench 测试 ---"
    
    mysql -h 127.0.0.1 -P "$port" -u"$MYSQL_USER" -p"$MYSQL_PWD" -e "CREATE DATABASE IF NOT EXISTS \`${db_name}\`;"
    if [ $? -ne 0 ]; then
        echo "错误: [端口 ${port}] 无法连接到 MySQL 或创建数据库。跳过此实例。"
        return 1
    fi

    if [ "$PREPARE_DATA" == "true" ]; then
        echo "[端口 ${port}] 正在准备数据 (TS=${TABLE_SIZE}, TBL=${TABLE_COUNT}, DB=${db_name})..."
        "$SYSBENCH_EXECUTABLE" "${LUA_SCRIPT_BASE_PATH}/oltp_common.lua" \
            --mysql-host=127.0.0.1 --mysql-port="$port" \
            --mysql-user="$MYSQL_USER" --mysql-password="$MYSQL_PWD" \
            --mysql-db="$db_name" \
            --table-size=$TABLE_SIZE --tables=$TABLE_COUNT \
            --threads=$THREADS_PREPARE prepare
        if [ $? -ne 0 ]; then
            echo "错误: [端口 ${port}] 数据准备失败。"
            return 1
        fi
    fi

    for scenario in "${SCENARIOS[@]}"; do
        local lua_script=""
        case "$scenario" in
            read_write)       lua_script="${LUA_SCRIPT_BASE_PATH}/oltp_read_write.lua" ;;
            point_select)     lua_script="${LUA_SCRIPT_BASE_PATH}/oltp_point_select.lua" ;;
            read_only)        lua_script="${LUA_SCRIPT_BASE_PATH}/oltp_read_only.lua" ;;
            write_only)       lua_script="${LUA_SCRIPT_BASE_PATH}/oltp_write_only.lua" ;;
            update_index)     lua_script="${LUA_SCRIPT_BASE_PATH}/oltp_update_index.lua" ;;
            update_non_index) lua_script="${LUA_SCRIPT_BASE_PATH}/oltp_update_non_index.lua" ;;
            *) echo "错误: [端口 ${port}] 不支持的测试场景 '$scenario'。跳过。"; continue ;;
        esac

        if [ ! -f "$lua_script" ]; then
            echo "错误: [端口 ${port}] 找不到场景脚本: $lua_script"; continue
        fi
        
        local result_log_name="${scenario}_ts${TABLE_SIZE}_tbl${TABLE_COUNT}_th${THREADS_RUN}_result.log"
        echo "[端口 ${port}] 运行场景: $scenario (Threads=${THREADS_RUN})"
        
        ### 健壮性改进: 将标准错误(2)重定向到标准输出(1)，然后一起写入日志文件 ###
        # 这样，sysbench的任何错误信息（如连接失败）都会被记录下来
        "$SYSBENCH_EXECUTABLE" "$lua_script" \
            --mysql-host=127.0.0.1 --mysql-port="$port" \
            --mysql-user="$MYSQL_USER" --mysql-password="$MYSQL_PWD" \
            --mysql-db="$db_name" \
            --table-size=$TABLE_SIZE --tables=$TABLE_COUNT \
            --threads=$THREADS_RUN --time=$TEST_DURATION \
            --report-interval=$REPORT_INTERVAL run > "${output_dir}/${result_log_name}" 2>&1

        if [ $? -ne 0 ]; then
            echo "错误: [端口 ${port}] 场景 '$scenario' 测试失败。详情请查看: ${output_dir}/${result_log_name}"
        else
            echo "[端口 ${port}] 场景 '$scenario' 测试完成。"
        fi
    done
    return 0
}

# --- 并行运行所有测试 ---
echo "--- 并行运行 Sysbench: 场景=${SCENARIOS_STR}, TS=${TABLE_SIZE}, TBL=${TABLE_COUNT}, TH=${THREADS_RUN} ---"
pids=()
for port in "${PORTS[@]}"; do
    run_sysbench_for_instance "$port" &
    pids+=($!)
done

status=0
for pid in "${pids[@]}"; do
    wait "$pid" || status=1
done

if [ $status -ne 0 ]; then
    echo "警告: 一个或多个 Sysbench 测试未成功完成。"
    exit 1
fi

echo "所有 Sysbench 基准测试已完成。"